import os
import re

import chromadb
from chromadb.utils import embedding_functions
from docx import Document
from openai import OpenAI


def load_employee_file(file_path):
    """
    1:加载word文件
    """
    docx = Document(file_path)
    all_text = []
    for para in docx.paragraphs:
        # print(para.text)
        # 去除首尾空格 ,数据清洗的过程
        clean_text = para.text.strip()
        # 去除特殊符号
        # clean_text = re.sub(r'[^\w\s]', '', clean_text)
        all_text.append(clean_text)
    return "\n".join(all_text)


def split_text(text, chunk_size=200):
    """
    2:文本切割:按照固定长度切割
    """
    chunks = []
    for i in range(0, len(text), chunk_size):
        chunks.append(text[i:i + chunk_size])

    return chunks


def text_embeddings(texts):
    """
    texts = ["王科宇","陈治文","张三"]
    embeddings = [[1024],[1024],[1024]]

    """
    import os
    from openai import OpenAI

    # 检查API密钥是否存在
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("错误：未设置DASHSCOPE_API_KEY环境变量，无法进行文本嵌入")
        
    client = OpenAI(
        api_key=api_key,  # 如果您没有配置环境变量，请在此处用您的API Key进行替换
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"  # 百炼服务的base_url
    )
    if isinstance(texts, str):
        texts = [texts]

    batch_size = 10
    embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        completion = client.embeddings.create(
            model="text-embedding-v4",
            input=batch,
            dimensions=1024,  # 指定向量维度（仅 text-embedding-v3及 text-embedding-v4支持该参数）
            encoding_format="float"
        )
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings


def store_in_chromadb(chunks):
    embeddings = text_embeddings(chunks)
    if len(chunks) != len(embeddings):
        print("文本块数量和嵌入向量数量不一致")
        return
    client = chromadb.PersistentClient(path="./chroma_db")
    collection = client.get_or_create_collection(
        name="employee_manual",
        embedding_function=None
    )
    collection.add(
        documents=chunks,
        embeddings=embeddings,
        ids=[f"chunk_{i}" for i in range(len(chunks))]
    )

    print(f"成功往向量数据库employee_manual,存入{len(embeddings)}个向量")

def query_chroma(query_text):
    client = chromadb.PersistentClient(path="./chroma_db")
    
    # 检查API密钥是否存在
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        raise ValueError("错误：未设置DASHSCOPE_API_KEY环境变量，无法进行文本检索")
        
    embedding_function= embedding_functions.OpenAIEmbeddingFunction(
        api_key=api_key,
        model_name="text-embedding-v4",
        api_base="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    collection = client.get_or_create_collection(
        name="employee_manual",
        embedding_function=embedding_function
    )
    results = collection.query(
        query_texts=[query_text],
        n_results=5,
        include=["documents", "distances"]
    )
    return results

def rag_answer_question(results, question):
    """
    rag:检索-聚合—生成
    """
    print("results:", results)
    # 检查是否有检索到文档
    if not results["documents"] or not results["documents"][0]:
        print("未检索到相关文档内容")
        return "未找到相关文档内容，无法回答该问题。"
    
    retrieved_chunks = list(set(results["documents"][0]))
    print(f"去重之后的结果列表：{retrieved_chunks}")
    context= "\n".join([f"- {chunk}" for chunk in retrieved_chunks])
    print(f"\n[检索到的参考内容]\n{context}\n")
    
    # 使用阿里云百炼平台的API密钥
    api_key = os.getenv("DASHSCOPE_API_KEY")
    if not api_key:
        print("错误：未设置DASHSCOPE_API_KEY环境变量")
        return "系统配置错误：未设置API密钥，无法生成回答。"
        
    try:
        ai_client = OpenAI(
            api_key=api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        )
    except Exception as e:
        print(f"创建AI客户端时出错: {e}")
        return "系统配置错误：无法创建AI客户端。"
        
    prompt = f"""
    你是一个专业的HR助理，请仔细阅读以下员工手册内容并准确回答用户问题
    [员工手册]
    {context}
    [用户问题]
    {question}
    [助手回答]:
    1.严格基于员工手册内容进行回答，请勿编造答案。
    2.如果找不到答案，请直接返回"无法回答"。
    3.请将你的回答限制在200字以内。
    4.如有多个要点，请分点列出，
    5.如果涉及具体条款，可引用相关内容
    """
    print(f"提示词：{prompt}")
    
    try:
        completion = ai_client.chat.completions.create(
            model="qwen-plus",
            messages=[
                {"role": "system", "content": prompt},
            ],
        )
        print(completion.choices[0].message.content)
        return completion.choices[0].message.content
    except Exception as e:
        print(f"调用大模型时出错: {e}")
        return "调用大模型时出错，请检查API密钥配置和网络连接。"

if __name__ == '__main__':
    import os
    
    # 检查必要的环境变量
    dashscope_api_key = os.getenv("DASHSCOPE_API_KEY")
    openai_api_key = os.getenv("OPENAI_API_KEY")
    
    if not dashscope_api_key:
        print("警告：未设置DASHSCOPE_API_KEY环境变量")
        print("请设置DASHSCOPE_API_KEY环境变量以使用文本嵌入和检索功能")
        
    if not openai_api_key:
        print("警告：未设置OPENAI_API_KEY环境变量")
        print("请设置OPENAI_API_KEY环境变量以使用大模型问答功能")
    
    # 确保在正确的目录下查找文件
    script_dir = os.path.dirname(os.path.abspath(__file__))
    manual_path = os.path.join(script_dir, "employee_manual.docx")
    
    # 检查员工手册文件是否存在
    if not os.path.exists(manual_path):
        print(f"错误：找不到员工手册文件 {manual_path}")
        exit(1)
        
    # 加载并处理员工手册
    all_text = load_employee_file(manual_path)
    chunks = split_text(all_text)
    print(f"文档已切割为 {len(chunks)} 个文本块")
    
    # 存储到向量数据库
    try:
        store_in_chromadb(chunks)
    except ValueError as e:
        print(f"存储向量时出错: {e}")
        exit(1)
    
    # 测试查询
    query_text = "员工上班迟到15分钟,怎么处理"
    try:
        results = query_chroma(query_text)
        if results:
            rag_answer_question(results, query_text)
    except ValueError as e:
        print(f"查询时出错: {e}")
        exit(1)













